/**
 * @file matmul_call.cpp
 *
 * Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include<stdio.h>
#include<iostream>
#include "data_utils.h"
#include "param.h"
#include "kernel_tiling/kernel_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "acl/acl_base.h"
#include "aclrtlaunch_matmul_custom.h"
#else
#include "tikicpulib.h"
extern "C" void matmul_custom(uint8_t *a, uint8_t *b, uint8_t *c, uint8_t *workspace, uint8_t *tiling);
#endif
extern void GenerateTiling(const char *socVersion, uint8_t *tilingBuf,int M,int N,int K,uint32_t blockDim);
template<typename T>
void  ValueCast(uint16_t* a,T* a_cast,int count){
  for(int i=0;i<count;++i){
    a_cast[i]=static_cast<T>(a[i]);
  }
}
template<typename T>
void ValueTurnBack(uint16_t* c,T* c_cast,int count){
    for(int i=0;i<count;++i){
        c[i]=static_cast<uint16_t>(c_cast[i]);
    }
}

extern "C" int32_t matmul_custom_call(uint16_t *a, uint16_t *b, uint16_t *c,size_t M,size_t N,size_t K)
{
    using half=aclFloat16;
    const char *socVersion = SOC_VERSION;
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(socVersion);

    size_t a_size = M * K * sizeof(float);
    size_t b_size = K * N * sizeof(float); 
    size_t c_size = M * N * sizeof(float);
    
    // matmul TCubeTiling + localMemorySize
    size_t tilingFileSize = sizeof(TCubeTiling) + sizeof(uint64_t);
    size_t userWorkspaceSize = 0;
    size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform->GetLibApiWorkSpaceSize());
    size_t workspaceSize = userWorkspaceSize + systemWorkspaceSize;
    uint8_t *tilingBuf = (uint8_t *)malloc(tilingFileSize);

#ifdef CUSTOM_ASCEND310P
    uint32_t blockDim = 2;
#else
//blockDim处修改
    uint32_t blockDim = 20;
#endif
    GenerateTiling(socVersion, tilingBuf,M,N,K,blockDim);
    printf("3\n");
#ifdef ASCENDC_CPU_DEBUG
    uint8_t *a = (uint8_t *)AscendC::GmAlloc(a_size);
    uint8_t *b = (uint8_t *)AscendC::GmAlloc(b_size);
    uint8_t *c = (uint8_t *)AscendC::GmAlloc(c_size);
    uint8_t *workspace = (uint8_t *)AscendC::GmAlloc(workspaceSize);
    uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tilingFileSize);
    memcpy_s(tiling, tilingFileSize, tilingBuf, tilingFileSize);

    ICPU_RUN_KF(matmul_custom, blockDim, a, b, c, workspace, tiling);

    WriteFile("./output/output.bin", c, c_size);

    AscendC::GmFree((void *)a);
    AscendC::GmFree((void *)b);
    AscendC::GmFree((void *)c);
    AscendC::GmFree((void *)workspace);
    AscendC::GmFree((void *)tiling);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));
    float* a_cast;
    float* b_cast;
    float* c_cast;
    CHECK_ACL(aclrtMallocHost((void**)&a_cast,a_size));
    CHECK_ACL(aclrtMallocHost((void**)&b_cast,b_size));
    CHECK_ACL(aclrtMallocHost((void**)&c_cast,c_size));
    uint8_t *aHost;
    uint8_t *aDevice;
    CHECK_ACL(aclrtMallocHost((void**)(&aHost),a_size));
    CHECK_ACL(aclrtMalloc((void **)&aDevice, a_size, ACL_MEM_MALLOC_HUGE_FIRST));
    ValueCast<float>(a,a_cast,M*K);
    CHECK_ACL(aclrtMemcpy(aHost,a_size,reinterpret_cast<uint8_t*>(a_cast),a_size,ACL_MEMCPY_HOST_TO_HOST));
    float* a_h=reinterpret_cast<float*>(aHost);
    CHECK_ACL(aclrtMemcpy(aDevice, a_size, aHost, a_size, ACL_MEMCPY_HOST_TO_DEVICE));
    uint8_t* bHost;
    uint8_t* bDevice;
    CHECK_ACL(aclrtMallocHost((void**)(&bHost),b_size));
    CHECK_ACL(aclrtMalloc((void **)&bDevice, b_size, ACL_MEM_MALLOC_HUGE_FIRST));
    ValueCast<float>(b,b_cast,K*N);
    CHECK_ACL(aclrtMemcpy(bHost,b_size,reinterpret_cast<uint8_t*>(b_cast),b_size,ACL_MEMCPY_HOST_TO_HOST));
    CHECK_ACL(aclrtMemcpy(bDevice, b_size, bHost, b_size, ACL_MEMCPY_HOST_TO_DEVICE));
    uint8_t *workspaceDevice;
    CHECK_ACL(aclrtMalloc((void **)&workspaceDevice, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));

    uint8_t *tilingHost;
    uint8_t *tilingDevice;
    CHECK_ACL(aclrtMallocHost((void **)(&tilingHost), tilingFileSize));
    CHECK_ACL(aclrtMalloc((void **)&tilingDevice, tilingFileSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemcpy(tilingHost, tilingFileSize, tilingBuf, tilingFileSize, ACL_MEMCPY_HOST_TO_HOST));
    CHECK_ACL(aclrtMemcpy(tilingDevice, tilingFileSize, tilingHost, tilingFileSize, ACL_MEMCPY_HOST_TO_DEVICE));
    
    uint8_t *cHost;
    uint8_t *cDevice;
     CHECK_ACL(aclrtMallocHost((void**)(&cHost),c_size));
    CHECK_ACL(aclrtMalloc((void **)&cDevice, c_size, ACL_MEM_MALLOC_HUGE_FIRST));
    
    ACLRT_LAUNCH_KERNEL(matmul_custom)
    (blockDim, stream, aDevice, bDevice, cDevice, workspaceDevice, tilingDevice);
    CHECK_ACL(aclrtSynchronizeStream(stream));
    CHECK_ACL(aclrtMemcpy(cHost , c_size, cDevice, c_size, ACL_MEMCPY_DEVICE_TO_HOST));
    CHECK_ACL(aclrtMemcpy(reinterpret_cast<uint8_t*>(c_cast),c_size,cHost,c_size,ACL_MEMCPY_HOST_TO_HOST));
    c_cast=reinterpret_cast<float*>(cHost);
    int iszero=0;
    for(int i=0;i<M*N;++i){
        if(c_cast[i]==0){
            iszero+=1;
        }
    }
    std::cout<<"N*M:"<<N*M<<" now it will judge c_cast's elements zero num: "<<iszero<<std::endl;
    ValueTurnBack<float>(c,c_cast,M*N);
    CHECK_ACL(aclrtFreeHost(aHost))
    CHECK_ACL(aclrtFree(aDevice));
    CHECK_ACL(aclrtFreeHost(bHost));
    CHECK_ACL(aclrtFree(bDevice));
    CHECK_ACL(aclrtFree(workspaceDevice));
    CHECK_ACL(aclrtFree(tilingDevice));
    CHECK_ACL(aclrtFreeHost(tilingHost));
    CHECK_ACL(aclrtFree(cDevice));
    CHECK_ACL(aclrtFreeHost(cHost));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    free(tilingBuf);
    return 0;
}